#!/usr/bin/python3

import roslib
import rospy
import math
import tf
import tf2_ros
from geometry_msgs.msg import Pose, Twist, TransformStamped, PointStamped, PoseStamped
import time
import numpy as np
from std_msgs.msg import Float64MultiArray
from sensor_msgs.msg import JointState
import copy
import sys
from tf2_geometry_msgs import tf2_geometry_msgs
import roslibpy
from digital_twining_bridge import twining_bridge
import cv2
import base64
from sensor_msgs.msg import CompressedImage, Image
from std_msgs.msg import Float64MultiArray
from cv_bridge import CvBridge, CvBridgeError
import json

class tail_twining_bridge(twining_bridge):
    def __init__(self,nodename):
        rospy.init_node(nodename)
        debug_prefix = "/p903_system/"
        # debug_prefix = ""

        if(rospy.has_param(debug_prefix + "digital_twin_server_config")):
            digital_twin_server_config = rospy.get_param(debug_prefix + 'digital_twin_server_config')
        else:
            digital_twin_server_config = {}
            digital_twin_server_config['ip_address'] = '127.0.0.1'
            digital_twin_server_config['interface'] = 'rosbridge'
            digital_twin_server_config['system']= 'noetic'
        
        self.rosServer = None
        self.joint_name_index = None
        self.realtime_tracking_switch_on = True

        self.rosServer = self.getValidRosServer(digital_twin_server_config['ip_address'])
        self.sim_clap_pub = roslibpy.Topic(self.rosServer,'/weiduan_targeting/clap_joint_states', 'sensor_msgs/JointState')
        self.sim_rotors2_pub = roslibpy.Topic(self.rosServer,'/zhongduan_rotors2/rotors2_joint_states', 'sensor_msgs/JointState')

        self.clap_joint_states = rospy.Subscriber("/weiduan_targeting/joint_states",JointState,self.weiduanTargetingJointStates_cb)
        self.rotors2_joint_states = rospy.Subscriber("/zhongduan_rotors2/joint_states",JointState,self.zhongduanRotors2JointStates_cb)
        # self.turn_on_realtime_tracking = rospy.Subscriber("turn_on_realtime_tracking",Bool, self.turnOnRealtimeTracking_cb)

        self.cvbridge = CvBridge()

        tele1 = roslibpy.Topic(self.rosServer, "/weiduan_targeting/telescope1/image_raw/compressed", "sensor_msgs/CompressedImage", queue_size=1, throttle_rate=100)
        tele2 = roslibpy.Topic(self.rosServer, "/weiduan_targeting/telescope2/image_raw/compressed", "sensor_msgs/CompressedImage", queue_size=1, throttle_rate=100)
        inspect1 = roslibpy.Topic(self.rosServer, "/weiduan_targeting/image_processor/inspect1", "sensor_msgs/CompressedImage", queue_size=1, throttle_rate=100)
        inspect2 = roslibpy.Topic(self.rosServer, "/weiduan_targeting/image_processor/inspect2", "sensor_msgs/CompressedImage", queue_size=1, throttle_rate=100)
        circle = roslibpy.Topic(self.rosServer, "/weiduan_targeting/image_processor/circle_result", "std_msgs/Float64MultiArray", queue_size=1, throttle_rate=100)

        self.update_parameter = roslibpy.Topic(self.rosServer, 
            "/weiduan_targeting/image_processor/process_parameter", 
            "std_msgs/Float64MultiArray", 
            queue_size=1, throttle_rate=100)

        tele1.subscribe(self.image1_cb)
        tele2.subscribe(self.image2_cb)
        inspect1.subscribe(self.inspect1_cb)
        inspect2.subscribe(self.inspect2_cb)
        circle.subscribe(self.circle_cb)

        self.dt_image1 = rospy.Publisher("/weiduan_targeting/digital_twin/telescope1/image_raw", Image,queue_size=1)
        self.dt_image2 = rospy.Publisher("/weiduan_targeting/digital_twin/telescope2/image_raw", Image,queue_size=1)
        self.dt_inspect1 = rospy.Publisher("/weiduan_targeting/digital_twin/image_processor/inspect1", Image,queue_size=1)
        self.dt_inspect2 = rospy.Publisher("/weiduan_targeting/digital_twin/image_processor/inspect2", Image,queue_size=1)
        self.dt_circle = rospy.Publisher("/weiduan_targeting/digital_twin/image_processor/circle_result", Float64MultiArray , queue_size=1)

        self.transfer_paremeter = rospy.Subscriber("/weiduan_targeting/digital_twin/image_processor/process_parameter", 
            Float64MultiArray,self.para_cb)

        self.tfBroadcaster=tf2_ros.StaticTransformBroadcaster()
        self.staticTransformList = []

        if(rospy.has_param("key_components")):
            self.component_names = rospy.get_param("key_components")
            self.key_components = {}
            for i in self.component_names:
                self.key_components[i] = rospy.get_param(i+"/born")
        else:
            self.key_components = {}

        self.establish_key_components(self.key_components)

    def weiduanTargetingJointStates_cb(self,msg):
        self.republish_by_roslibpy(msg,self.sim_clap_pub)
        return

    def zhongduanRotors2JointStates_cb(self,msg):
        self.republish_by_roslibpy(msg,self.sim_rotors2_pub)
        return

    def republish_by_roslibpy(self,msg,pub):
        # if(not self.realtime_tracking_switch_on):
        #     return 

        # if(self.joint_name_index is None):
        #     para_name = "/pogo_group/ptp_controller/joints"
        #     if(rospy.has_param(para_name)):
        #         self.joint_names = rospy.get_param(para_name)
        #         self.joint_name_index = {self.joint_names[i]:i for i in range(len(self.joint_names))}
        #     else:
        #         return

        # sent_command=[0.0]*len(self.joint_names)
        # for i in range(len(msg.name)):
        #     sent_command[self.joint_name_index[msg.name[i]]] = msg.position[i]
        sent_command=roslibpy.Message({
            'header':roslibpy.Header(0,roslibpy.Time(0,0),''),
            'name':msg.name,
            'position':msg.position
        })

        pub.publish(sent_command)
        return

    def get_colored_image_msg(self,msg):
        base64_bytes = msg['data'].encode('ascii')
        image_bytes = base64.b64decode(base64_bytes)
        # cmsg = CompressedImage()
        # cmsg.format = msg['format']
        # cmsg.data = msg['data'].encode('ascii')
        np_arr = np.frombuffer(image_bytes, np.uint8)
        cv_img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

        # bridge = CvBridge()

        # cv_img = bridge.compressed_imgmsg_to_cv2(cmsg)
        pub_msg = self.cvbridge.cv2_to_imgmsg(cv_img,encoding='bgr8')
        return pub_msg

    def get_mono_image_msg(self,msg):
        base64_bytes = msg['data'].encode('ascii')
        image_bytes = base64.b64decode(base64_bytes)
        # cmsg = CompressedImage()
        # cmsg.format = msg['format']
        # cmsg.data = msg['data'].encode('ascii')
        np_arr = np.frombuffer(image_bytes, np.uint8)
        cv_img = cv2.imdecode(np_arr, cv2.IMREAD_GRAYSCALE)

        # bridge = CvBridge()

        # cv_img = bridge.compressed_imgmsg_to_cv2(cmsg)
        pub_msg = self.cvbridge.cv2_to_imgmsg(cv_img,encoding='mono8')
        return pub_msg

    def image1_cb(self,msg):
        pub_msg = self.get_colored_image_msg(msg)
        self.dt_image1.publish(pub_msg)

    def image2_cb(self,msg):
        pub_msg = self.get_colored_image_msg(msg)
        self.dt_image2.publish(pub_msg)

    def inspect1_cb(self,msg):
        pub_msg = self.get_mono_image_msg(msg)
        self.dt_inspect1.publish(pub_msg)

    def inspect2_cb(self,msg):
        pub_msg = self.get_mono_image_msg(msg)
        self.dt_inspect2.publish(pub_msg)

    def circle_cb(self,msg):
        pub_msg = Float64MultiArray()
        pub_msg.data = msg['data']
        self.dt_circle.publish(pub_msg)

    def para_cb(self,msg):
        pub_msg=roslibpy.Message()
        pub_msg['data']=msg.data
        self.update_parameter.publish(pub_msg)

    # def turnOnRealtimeTracking_cb(self,msg):
    #     print(msg)
    #     return

    def run(self):
        '''
        rate = rospy.Rate(50)
        trans = TransformStamped()
        while not rospy.is_shutdown():
            if(self.msgNames is None or self.msgPose is None):
                rate.sleep()
                continue

            for i in range(len(self.msgNames)):
                if(self.msgNames[i] in self.baseLinks.keys()):
                    #print(msg.name[i])
                    #print(msg.pose[i])
                    trans.header.frame_id = "world"
                    trans.child_frame_id = self.baseLinks[self.msgNames[i]]

                    trans.header.stamp = rospy.Time.now()

                    trans.transform.translation.x = self.msgPose[i].position.x
                    trans.transform.translation.y = self.msgPose[i].position.y
                    trans.transform.translation.z = self.msgPose[i].position.z

                    trans.transform.rotation.x = self.msgPose[i].orientation.x
                    trans.transform.rotation.y = self.msgPose[i].orientation.y
                    trans.transform.rotation.z = self.msgPose[i].orientation.z
                    trans.transform.rotation.w = self.msgPose[i].orientation.w

                    self.tfBroadcaster.sendTransform(trans)
            rate.sleep()
        '''
        rospy.spin()

if __name__ == '__main__':

    r = tail_twining_bridge('rotors2_twining_bridge')
    r.run()
